
from .sessionInfo import SessionInfo
from ..execute.baseExeHis import BaseExeHis
from abc import ABC, abstractmethod
class Msg:
    def __init__(self, content: str, role: str, msg_type: str, **kwargs): 
        '''
            role is [user, assistant]
            msg_type is [text, audio, image]
        '''
        self.content = content # image_url if type is image
        self.role = role
        self.type = msg_type
        self.kwargs = kwargs

    def msg_dict(self, **kwargs) -> dict:
        if self.type == "text":
            return {"content": self.content, "role": self.role, "kwargs":kwargs}
        elif self.type == "image":
            return {"role": kwargs.get('role', "user"), "content":[
                                                    {"type": "text", "text": "The image is:"},
                                                    {   "type": "image_url",
                                                        "image_url": { "url": self.content }
                                                    }
                                                ]}
                                            

class BaseGraphNode(ABC):
    def __init__(self, id:str, data:dict, graphNodeManager ,random_key:str):
        from managers.graphNodeManager import GraphNodeManager
        self.id = id
        self.data = data
        assert isinstance(graphNodeManager, GraphNodeManager), "graphNodeManager must be an instance of GraphNodeManager"
        self.graphNodeManager = graphNodeManager
        self.random_key = random_key
        self.name = data['data']['name']
        self.type = data['type']
        self.dict = {'in':{}, 'out':{}} #'out' :{sessionInfo.id:{ targethandle:[output_1, output_2] }}
        self.executed_num = 0
        self.flow_next_nodes = []
        
    def get_new_exe_his(self, session_info_ID, beforerunning=False):
        return BaseExeHis(self.graphNodeManager, self.graphNodeManager.project.id, session_info_ID, self.id, beforerunning=beforerunning)

    def execute(self, session_info: SessionInfo=None, exe_his:BaseExeHis=None, trigger=None, running_id=None, **kwargs):
        pass

    # 将默认实现直接放在基类中
    def add_flow_next(self, node_id:str, **kwargs):
        """添加 flow_next_nodes。
        如果重写此方法，必须同时重写 get_all_flow_next_nodes。
        """
        self.flow_next_nodes.append(node_id)

    def get_all_flow_next_nodes(self): 
        """获取所有 flow_next_nodes。
        如果重写此方法，必须同时重写 add_flow_next。
        """
        return self.flow_next_nodes

    @classmethod
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        has_add = 'add_flow_next' in cls.__dict__
        has_get = 'get_all_flow_next_nodes' in cls.__dict__
        if has_add != has_get:
            raise TypeError(
                f"Class {cls.__name__} must override both add_flow_next and "
                f"get_all_flow_next_nodes methods together, or none of them."
            )
    
    def flow_next(self, session_info: SessionInfo, trigger, running_id, **kwargs):
        """将下一个节点添加到执行队列中"""
        for node_id in self.flow_next_nodes:
            # 更新执行次数
            session_info.flow_prev_nodes_executed_number[node_id][self.id] += 1
            self.graphNodeManager.add_to_execution_queue(node_id, session_info,trigger=trigger, running_id=running_id, **kwargs)
            
    def flow_next_process(self, session_info: SessionInfo, nodes:list[str], trigger, running_id, **kwargs):
        session_info.done_exe_with_nodeID(self.id)
        if trigger and trigger.enInterrupt and trigger.session_id_2_running_id[session_info.id]!=running_id:return
        for node in nodes:
            session_info.flow_prev_nodes_executed_number[node][self.id] += 1
            self.graphNodeManager.add_to_execution_queue(node, session_info,trigger=trigger, running_id=running_id, **kwargs)

    def get_newest_output_in_session(self, session_info: SessionInfo, target_handle: str):
        with session_info.out_write_lock:
            if session_info.out.get(self.id, None) is None:return None

            if session_info.out[self.id].get(target_handle, None) is None:return None

            if session_info.out[self.id][target_handle]:
                return session_info.out[self.id][target_handle][-1]
        return None

    def stop(self):
        pass

    def clear(self, session_info: SessionInfo):
        with session_info.out_write_lock:
            if session_info.out.get(self.id, None):
                session_info.out[self.id].clear()